import argparse
import os
import shutil
import time

import torch
import yaml
from tqdm import tqdm
from models import transfer
from utils.datasets import create_dataloader, LoadImages
from utils.torch_utils import cuda2cpu, select_device


def predict(opt):
    if os.path.exists(opt.output):
        shutil.rmtree(opt.output)  # delete output folder
    os.makedirs(opt.output)  # make new output folder

    device = select_device(opt.device)

    with open("data/abnormal.yaml",encoding="utf-8") as f:
        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # model dict

    nc = len(data_dict["names"])
    # model = transfer.resnet18(nc=nc,pretrained=False)
    model = transfer.Model(opt.model_name, nc=nc, pretrained=True).to(device)  # create
    model.load_state_dict(torch.load(opt.weights, map_location=device)["model"], strict=False)
    model.eval()
    model.to(device)

    predict_dataset = LoadImages(opt.source, img_size=opt.img_size)

    predtimes = []
    pbar = tqdm(enumerate(predict_dataset), total=len(predict_dataset))
    for i, (input, path) in pbar:
        input = input.to(device)

        if input.ndimension() == 3:
            input = input.unsqueeze(0)

        start = time.time()
        output = model(input)
        prob, pred = torch.max(output, 1)

        pred_array = cuda2cpu(pred).tolist()
        predvalue = pred_array[0]
        end = time.time()

        classnameDir = os.path.join(opt.output, data_dict["names"][predvalue])
        if not os.path.exists(classnameDir):
            os.makedirs(classnameDir)

        shutil.copy(path, classnameDir)

        end_start_time = end - start
        predtimes.append(end_start_time)
        s = "{} Predict a picture,time:{}".format(device, end_start_time)
        pbar.set_description(s)
    print("{}平均预测时间：{}".format(device, sum(predtimes) / len(predtimes)))
    print('*' * 50)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model-name', type=str, default='shufflenet_v2_x0_5', help='initial weights path')  # 要使用那个模型来训练
    parser.add_argument('--weights', nargs='+', type=str, default='runs/exp9_abnormal_shufflenet_v2_x0_5/weights/best.pt',
                        help='model.pt path(s)')
    parser.add_argument('--source', type=str,
                        default=r'\\10.20.200.170\chintAI\ext\scene\2020-10-8', #2020-10-7
                        help='source')  # file/folder, 0 for webcam
    parser.add_argument('--output', type=str, default='inference/output', help='output folder')  # output folder
    parser.add_argument('--img-size', type=int, default=[16, 16], help='inference size (pixels)')
    parser.add_argument('--device', default='1', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    opt = parser.parse_args()

    with torch.no_grad():
        predict(opt)
